from typing import Dict, Any, Optional
from pydantic import BaseModel, Field, ValidationError
from langgraph.graph import StateGraph

class IdentityScanData(BaseModel):
    document_type: str = Field(description="证件类型：passport 或 id_card")
    mrz: Optional[str] = None
    id_no: Optional[str] = None

class ServiceIdentity(BaseModel):
    type: str = Field(default="identity")
    verified: bool
    reason: Optional[str] = None

def build_identity_subgraph():
    g = StateGraph(dict)

    def verify(state: Dict[str, Any]) -> Dict[str, Any]:
        raw_scan = state.get("scan_data") or {}
        try:
            scan = IdentityScanData.model_validate(raw_scan)
        except ValidationError:
            scan = IdentityScanData(document_type="unknown")
        doc_type = scan.document_type
        mrz = scan.mrz
        verified = False
        reason = None
        if doc_type == "passport":
            if mrz and isinstance(mrz, str):
                m = mrz.strip().splitlines()
                l1 = m[0].strip() if len(m) > 0 else ""
                l2 = m[1].strip() if len(m) > 1 else ""
                if len(l1) == 44 and len(l2) == 44:
                    verified = True
                else:
                    reason = "invalid_mrz"
            else:
                reason = "invalid_mrz"
        elif doc_type == "id_card":
            id_no = scan.id_no
            if id_no and len(str(id_no)) >= 10:
                verified = True
            else:
                reason = "invalid_id"
        else:
            reason = "unknown_document"
        service = {"type": "identity", "verified": verified}
        if reason:
            service["reason"] = reason
        try:
            svc = ServiceIdentity.model_validate(service).model_dump()
        except ValidationError:
            svc = {"type": "identity", "verified": False, "reason": "schema_error"}
        return {**state, "service": svc}

    g.add_node("verify", verify)
    g.set_entry_point("verify")
    return g.compile()